package utils

import (
	"math/rand"
)

type inlineRandom struct {
	Key     interface{}
	Value   int
	current int
}

type RandomFactory interface {
	Add(key interface{}, value int)
	Size() int
	Random() interface{}
}

// 按比率进行概率随机
// 比率越大概率越大，相同比率，相同概率
type RatioRandom struct {
	items  []inlineRandom
	all    int
	init   bool
	result []int
}

func NewRatioRandom() *RatioRandom {
	rr := &RatioRandom{
		items: make([]inlineRandom, 0),
		all:   0,
		init:  false,
	}
	return rr
}

func (rr *RatioRandom) Add(key interface{}, ratio int) {
	rr.items = append(rr.items, inlineRandom{
		key,
		ratio,
		0,
	})
}

func (rr *RatioRandom) Size() int {
	return rr.all
}

func (rr *RatioRandom) Compute() {
	rr.all = 0
	for key := range rr.items {
		rr.items[key].current = rr.items[key].Value
		rr.all += rr.items[key].Value
	}
	rr.result = make([]int, rr.all)
	index := 0
	for key := range rr.items {
		max := index + rr.items[key].current
		for i := index; i < max; i++ {
			rr.result[i] = key
		}
		index = max
	}

	rr.init = true
}

func (rr *RatioRandom) Random() interface{} {
	if !rr.init {
		rr.Compute()
	}
	r := rand.Int31n(int32(rr.all))
	index := rr.result[r]
	return rr.items[index].Key
}

// 按权重进行概率随机
// 权重越大概率越小，相同权重，相同概率
type WeightRandom struct {
	items     []inlineRandom
	maxWeight int
	all       int
	init      bool
	result    []int
}

func NewWeightRandom() *WeightRandom {
	wr := &WeightRandom{
		items:     make([]inlineRandom, 0),
		maxWeight: 0,
		all:       0,
		init:      false,
	}
	return wr
}

func (wr *WeightRandom) Add(key interface{}, weight int) {
	if weight >= wr.maxWeight {
		wr.maxWeight = weight + 1
	}
	wr.items = append(wr.items, inlineRandom{
		key,
		weight,
		0,
	})
}

func (wr *WeightRandom) Size() int {
	return wr.all
}

func (wr *WeightRandom) Compute() {
	wr.all = 0
	for key := range wr.items {
		current := wr.maxWeight - wr.items[key].Value
		wr.items[key].current = current
		wr.all += current
	}
	wr.result = make([]int, wr.all)
	index := 0
	for key := range wr.items {
		max := index + wr.items[key].current
		for i := index; i < max; i++ {
			wr.result[i] = key
		}
		index = max
	}

	wr.init = true
}

func (wr *WeightRandom) Random() interface{} {
	if !wr.init {
		wr.Compute()
	}
	r := rand.Int31n(int32(wr.all))
	index := wr.result[r]
	return wr.items[index].Key
}
